from datetime import datetime, timezone

from app.config.database import sql_alchemy_data_layer
from app.config.logging import logger
from app.core.chat_engine import ChatEngine
from app.models.chat_payload import ChatPayload
from chainlit_starter import *
from chainlit_mcp import *


## 将chainlit挂载到FastAPI的子应用

@cl.on_message
async def on_message(message: cl.Message):
    current_session_id = cl.context.session.id
    logger.info(f"session_id: {current_session_id} input message: {message.content}")
    payload = ChatPayload(session_id=current_session_id, prompt=message.content,
                          parent_message_id=message.parent_id)
    # 定义返回值
    # 获取当前UTC时间（精确到微秒）
    now = datetime.now(timezone.utc)
    # 格式化为ISO 8601并强制使用'Z'时区标记
    formatted_time = now.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
    # 必须的设置created_at,否则不会显示
    msg = cl.Message(content="", created_at=formatted_time)
    chat_engine = ChatEngine(payload)
    async for event in chat_engine.astream_events():
        if event.event_type == "ready":
            # 开始处理
            logger.info(f"chat ready finish.session_id:{current_session_id}")
            continue
        elif event.event_type == "token":
            # 输出结果
            await msg.stream_token(event.format_data())
        elif event.event_type == "complete":
            await msg.stream_token("")
        else:
            # 其他event_type 统一认为是异常event
            logger.info(f"chat error.session_id:{current_session_id}")
            await msg.stream_token(event.format_data())

    await msg.update()


@cl.on_stop
async def on_stop():
    """
    stop停止生成Event
    :return:
    """


@cl.on_chat_end
async def on_chat_end():
    """
    聊天结束/异常终止
    :return:
    """


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    # Fetch the user matching username from your database
    # and compare the hashed password with the value stored in the database
    if (username, password) == ("admin", "admin"):
        return cl.User(
            identifier="admin", metadata={"role": "admin", "provider": "credentials"}
        )
    else:
        return None


@cl.data_layer
def get_data_layer():
    return sql_alchemy_data_layer
